#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2025-07-30


import threading
from queue import Queue
import numpy as np
import time
import os
from typing import Union
from spirems import Subscriber, Publisher, cvimg2sms, sms2cvimg, def_msg, QoS, BaseNode, get_extra_args
from spirems.mod_helper import download_model
import argparse
import platform
import subprocess


class Wts2EngineNode_Trt(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=False,
            **kwargs
        )
        self.wts_model = self.get_param("wts_model", "sms::air2air_det_db-yolo11n_i640_c2-20250730.wts")
        self.engine_path = self.get_param("engine_path", "/home/nvidia/air2air_det_db-yolo11n_i640_c2-20250730.engine")
        self.task = self.get_param("task", "det")
        self.level = self.get_param("level", "n")
        self.imgsz = self.get_param("imgsz", 640)
        self.nc = self.get_param("nc", "2")
        self.nb = self.get_param("nb", 1)
        self.params_help()

        if self.wts_model.startswith("sms::"):
            self.local_wts_model = download_model(self.__class__.__name__, self.wts_model)
            assert self.local_wts_model is not None
        else:
            self.local_wts_model = self.wts_model

        current_path = os.path.dirname(os.path.abspath(__file__))
        program_fn = os.path.join(current_path, "build/yolo11_" + self.task)
        subprocess.run([program_fn, '-s', self.local_wts_model, self.engine_path, self.level, str(self.nc), str(self.imgsz), str(self.nb)])

        self.release()

    def release(self):
        BaseNode.release(self)

    def run(self):
        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name', '-j',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = Wts2EngineNode_Trt(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.join()
